// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package com.intellij.credentialStore.keePass

import com.intellij.credentialStore.*
import com.intellij.openapi.util.SystemInfo
import com.intellij.openapi.util.io.BufferExposingByteArrayOutputStream
import com.intellij.util.io.delete
import com.intellij.util.io.readBytes
import com.intellij.util.io.safeOutputStream
import org.yaml.snakeyaml.composer.Composer
import org.yaml.snakeyaml.nodes.*
import org.yaml.snakeyaml.parser.ParserImpl
import org.yaml.snakeyaml.reader.StreamReader
import org.yaml.snakeyaml.resolver.Resolver
import java.nio.file.NoSuchFileException
import java.nio.file.Path
import java.util.*

internal const val MASTER_KEY_FILE_NAME = "c.pwd"
private const val OLD_MASTER_PASSWORD_FILE_NAME = "pdb.pwd"

internal class MasterKey(value: ByteArray,
                         // is password auto-generated, used to force set master password if database file location changed
                         val isAutoGenerated: Boolean,
                         val encryptionSpec: EncryptionSpec) {
  var value: ByteArray? = value

  /**
   * Clear byte array to avoid sensitive data in memory
   */
  fun clear() {
    value?.fill(0)
    value = null
  }
}

internal class MasterKeyFileStorage(val passwordFile: Path) {
  fun load(): ByteArray? {
    var data: ByteArray
    var isOld = false
    try {
      data = passwordFile.readBytes()
    }
    catch (e: NoSuchFileException) {
      try {
        data = passwordFile.parent.resolve(OLD_MASTER_PASSWORD_FILE_NAME).readBytes()
      }
      catch (e: NoSuchFileException) {
        return null
      }
      isOld = true
    }

    try {
      val decrypted: ByteArray
      if (isOld) {
        decrypted = createBuiltInOrCrypt32EncryptionSupport(SystemInfo.isWindows).decrypt(data)
        data.fill(0)
      }
      else {
        decrypted = decryptMasterKey(data) ?: return null
      }
      data.fill(0)
      return decrypted
    }
    catch (e: Exception) {
      LOG.warn("Cannot decrypt master key, file content:\n${if (isOld) Base64.getEncoder().encodeToString(data) else data.toString(Charsets.UTF_8)}", e)
      return null
    }
  }

  private fun decryptMasterKey(data: ByteArray): ByteArray? {
    var encryptionType: EncryptionType? = null
    var value: ByteArray? = null
    for (node in createMasterKeyReader(data)) {
      val keyNode = node.keyNode
      val valueNode = node.valueNode
      if (keyNode is ScalarNode && valueNode is ScalarNode) {
        val propertyValue = valueNode.value ?: continue
        when (keyNode.value) {
          "encryption" -> encryptionType = EncryptionType.valueOf(propertyValue.toUpperCase())
          "value" -> value = Base64.getDecoder().decode(propertyValue)
        }
      }
    }

    if (encryptionType == null) {
      LOG.error("encryption type not specified in $passwordFile, default one will be used (file content:\n${data.toString(Charsets.UTF_8)})")
      encryptionType = getDefaultEncryptionType()
    }
    if (value == null) {
      LOG.error("password not specified in $passwordFile, automatically generated will be used (file content:\n${data.toString(Charsets.UTF_8)})")
      return null
    }

    // PGP key id is not stored in master key file because PGP encoded data already contains recipient id and no need to explicitly specify it,
    // so, this EncryptionSupport MUST be used only for decryption since pgpKeyId is set to null.
    return createEncryptionSupport(EncryptionSpec(encryptionType, pgpKeyId = null)).decrypt(value)
  }

  // passed key will be cleared
  fun save(key: MasterKey?) {
    if (key == null) {
      passwordFile.delete()
      return
    }

    val encodedValue = Base64.getEncoder().encode(createEncryptionSupport(key.encryptionSpec).encrypt(key.value!!))
    key.clear()

    val out = BufferExposingByteArrayOutputStream()
    val encryptionType = key.encryptionSpec.type
    out.writer().use {
      it.append("encryption: ").append(encryptionType.name).append('\n')
      it.append("isAutoGenerated: ").append(key.isAutoGenerated.toString()).append('\n')
      it.append("value: !!binary ")
    }

    passwordFile.safeOutputStream().use {
      it.write(out.internalBuffer, 0, out.size())
      it.write(encodedValue)
      encodedValue.fill(0)
    }
  }

  private fun readMasterKeyIsAutoGeneratedMetadata(data: ByteArray): Boolean {
    var isAutoGenerated = true
    for (node in createMasterKeyReader(data)) {
      val keyNode = node.keyNode
      val valueNode = node.valueNode
      if (keyNode is ScalarNode && valueNode is ScalarNode) {
        val propertyValue = valueNode.value ?: continue
        when (keyNode.value) {
          "isAutoGenerated" -> isAutoGenerated = propertyValue.toBoolean() || propertyValue == "yes"
        }
      }
    }
    return isAutoGenerated
  }

  fun isAutoGenerated(): Boolean {
    try {
      return readMasterKeyIsAutoGeneratedMetadata(passwordFile.readBytes())
    }
    catch (e: NoSuchFileException) {
      // true because on save will be generated
      return true
    }
  }
}

private fun createMasterKeyReader(data: ByteArray): List<NodeTuple> {
  val composer = Composer(ParserImpl(StreamReader(data.inputStream().reader())), object : Resolver() {
    override fun resolve(kind: NodeId, value: String?, implicit: Boolean): Tag {
      return when (kind) {
        NodeId.scalar -> Tag.STR
        else -> super.resolve(kind, value, implicit)
      }
    }
  })
  return (composer.singleNode as? MappingNode)?.value ?: emptyList()
}